/*  This file is part of Basepairs, an application for calculations on
    DNA Sequence polymorphism.
    Copyright (C) 2013 Kim Lindgren

    This program is free software: you can redistribute it and/or modify
    it under the terms of the GNU General Public License as published by
    the Free Software Foundation, either version 3 of the License, or
    (at your option) any later version.

    This program is distributed in the hope that it will be useful,
    but WITHOUT ANY WARRANTY; without even the implied warranty of
    MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
    GNU General Public License for more details.

    You should have received a copy of the GNU General Public License
    along with this program.  If not, see <http://www.gnu.org/licenses/>.*/

package basepairs;

import java.util.ArrayList;

public class Watterson {
    Watterson() {
        
    }
    
    public double Theta;
    public double synTheta;
    public double nonSynTheta;
    
    public double calcSegregatingSites(){ //original function, calculates for whole sequences. Gives same result as DnaSP.
        double pSeq;
        int mutations = 0;
        String first;
        String current;
        int seqs = Sequences.selected.size();
        ArrayList<String> variants = null;
        boolean mutated;
        
        for(int x = 0; x < Sequences.sequenceLength; x++){
            if(!Sequences.gaps.contains(x)){
                variants = new ArrayList<String>();
                mutated = false;
                
                first = Sequences.selected.get(0).substring(x,x+1);
                for(int y = 0; y < seqs; y++){
                    current = Sequences.selected.get(y).substring(x,x+1);
                    if(!current.equals(first)){
                        if(mutated == false){
                            mutated = true;
                            variants.add(current);
                        }
                        else if(!variants.contains(current)) {
                            variants.add(current);
                        }
                    }
                }
                
                mutations = mutations + variants.size();
            }
        }
        
        pSeq = mutations;
        double W = pSeq/(Sequences.sequenceLength-Sequences.gaps.size());
        this.Theta = W/Utility.calcHarmonic(seqs);
        
        return this.Theta;
    }
    
    public void calcSegregatingCodons(int firstOffset){ //modified function, works one codon at a time. Gives similar results to DnaSP.
        int synMutations = 0;
        int nonSynMutations = 0;
        String first;
        String current;
        int seqs = Sequences.selected.size();
        double length = Sequences.selected.get(0).length();
        double synLength = 0;
        double nonSynLength = 0;
        ArrayList<String> synVariants;
        ArrayList<String> nonSynVariants;
        
        boolean mutated;
        boolean synon;
                
        int sequenceLength = Sequences.separateCodons(Sequences.selected.get(0), firstOffset).size();
        
        System.out.println("Total number of codons " + sequenceLength);
                
        ArrayList<Integer> codonGaps = Sequences.findCodonsWithGaps(firstOffset);
                
        for(int codon = 0; codon < sequenceLength; codon++){
            if(!codonGaps.contains(codon)){
                synVariants = new ArrayList<String>();
                nonSynVariants = new ArrayList<String>();
                
                mutated = false;

                first = Sequences.separateCodons(Sequences.selected.get(0), firstOffset).get(codon);
                
                for(int sequence = 1; sequence < seqs; sequence++){
                    current = Sequences.separateCodons(Sequences.selected.get(sequence), firstOffset).get(codon);
                    
                    if(!first.equals(current)){
                        synon = Codon.isSynonymous(first, current);

                        for(int codonPosition = 0; codonPosition < 3; codonPosition++){
                            String s1 = first.substring(codonPosition,codonPosition+1);
                            String s2 = current.substring(codonPosition,codonPosition+1);

                            if(synon == true){
                                if(!s1.equals(s2)){
                                    if(mutated == false){
                                        mutated = true;
                                        synVariants.add(codonPosition + s2);
                                    }
                                    else if(!synVariants.contains(codonPosition + current)) {
                                        synVariants.add(codonPosition + s2);
                                    }
                                    
                                    synMutations++;
                                }
                                
                                synLength++;
                            }
                            else{
                                if(!s1.equals(s2)){
                                    if(mutated == false){
                                        mutated = true;
                                        nonSynVariants.add(codonPosition + s2);
                                    }
                                    else if(!nonSynVariants.contains(codonPosition + s2)) {
                                        nonSynVariants.add(codonPosition + s2);
                                    }
                                
                                    nonSynMutations++;
                                }
                                
                                nonSynLength++;
                            }
                        }
                    }
                    else{
                        synLength = synLength + 3;
                        nonSynLength = nonSynLength + 3;
                    }
                    
                    synMutations = synMutations + synVariants.size();
                    nonSynMutations = nonSynMutations + nonSynVariants.size();
                }
            }
        }
                
        double harmonic = Utility.calcHarmonic(seqs);
        
        double synW = (double)synMutations/(synLength);
        synTheta = synW/harmonic;
        
        double nonSynW = (double)nonSynMutations/(nonSynLength);
        nonSynTheta = nonSynW/harmonic;
    }
}
